type allocation = Frame.register Temp.Table.t

(* Doubly-linked list rep *)

type 'a dll = 'a Dllist.node_t option ref

type 'a dll_node = {
  mutable node : 'a Dllist.node_t option;
  mutable dll : 'a dll;
}

let create_dll () = ref None
let dll2list dll = Option.map_default Dllist.to_list [] !dll

let move_node setnode dll =
  let olddll = setnode.dll in
  let head = Option.get !olddll in
  let node = Option.get setnode.node in
  let next = Dllist.drop node in
  setnode.dll <- dll;
  (* Remove from old dll *)
  if next == node then olddll := None
  else if node == head then olddll := Some next;
  (* Append to new dll *)
  match !dll with
  | None -> dll := Some node
  | Some other -> Dllist.splice (Dllist.prev other) node

(* Node rep *)
type cnode = { inode : Igraph.node; n : int; setnode : cnode dll_node }

(* Move rep *)
type move = { src : cnode; dst : cnode; setnode : move dll_node }

module Cset = Set.Make (struct
  type t = cnode

  let compare = compare
end)

module Rset = Set.Make (struct
  type t = Frame.register

  let compare = compare
end)

let color ({ Liveness.graph; tnode; gtemp; moves } as igraph) spill_cost
    allocation registers =
  (* Nodeset functions *)
  let nextn = ref 0 in
  let add_node inode dll =
    let cnode = { inode; n = !nextn; setnode = { node = None; dll } } in
    nextn := !nextn + 1;
    let dllnode = Dllist.create cnode in
    cnode.setnode.node <- Some dllnode;
    (match !dll with
    | None -> dll := Some dllnode
    | Some other -> Dllist.splice (Dllist.prev other) dllnode);
    cnode
  in

  (* Moveset functions *)
  let add_move src dst dll =
    let move = { src; dst; setnode = { node = None; dll } } in
    let dllnode = Dllist.create move in
    move.setnode.node <- Some dllnode;
    (match !dll with
    | None -> dll := Some dllnode
    | Some other -> Dllist.splice (Dllist.prev other) dllnode);
    move
  in

  (* Start *)
  let n = List.length (Igraph.nodes graph) in
  let k = List.length registers in

  (* Node sets *)
  let precolored = create_dll () in
  let initial = create_dll () in
  let simplify_worklist = create_dll () in
  let freeze_worklist = create_dll () in
  let spill_worklist = create_dll () in
  let spilled_nodes = create_dll () in
  let coalesced_nodes = create_dll () in
  let colored_nodes = create_dll () in
  let select_stack = create_dll () in

  (* Move sets *)
  let coalesced_moves = create_dll () in
  let constrained_moves = create_dll () in
  let frozen_moves = create_dll () in
  let worklist_moves = create_dll () in
  let active_moves = create_dll () in

  (* Other data structures *)
  let adj_set = Array.make_matrix n n false in
  let adj_list = Array.make n Cset.empty in
  let degree = Array.make n 0 in
  let move_list = Array.make n [] in
  let alias = Array.make n None in
  let color = Array.make n None in

  let add_edge u v =
    if (not adj_set.(u.n).(v.n)) && u != v then (
      adj_set.(u.n).(v.n) <- true;
      adj_set.(v.n).(u.n) <- true;
      if u.setnode.dll != precolored then (
        adj_list.(u.n) <- Cset.add v adj_list.(u.n);
        degree.(u.n) <- degree.(u.n) + 1);
      if v.setnode.dll != precolored then (
        adj_list.(v.n) <- Cset.add u adj_list.(v.n);
        degree.(v.n) <- degree.(v.n) + 1))
  in

  let build () =
    let inodetab =
      List.fold_left
        (fun tab inode ->
          let temp = gtemp inode in
          let cnode =
            if Temp.Table.mem temp allocation then (
              let cnode = add_node inode precolored in
              color.(cnode.n) <- Some (Temp.Table.find temp allocation);
              cnode)
            else add_node inode initial
          in
          Igraph.Table.add inode cnode tab)
        Igraph.Table.empty (Igraph.nodes graph)
    in
    List.iter
      (fun inode ->
        let u = Igraph.Table.find inode inodetab in
        List.iter
          (fun adj -> add_edge u (Igraph.Table.find adj inodetab))
          (Igraph.adj inode))
      (Igraph.nodes graph);
    List.iter
      (fun (u, v) ->
        let u = Igraph.Table.find u inodetab in
        let v = Igraph.Table.find v inodetab in
        let m = add_move u v worklist_moves in
        move_list.(u.n) <- move_list.(u.n) @ [ m ];
        move_list.(v.n) <- move_list.(v.n) @ [ m ])
      moves
  in

  let node_moves n =
    List.filter
      (fun m ->
        m.setnode.dll == active_moves || m.setnode.dll == worklist_moves)
      move_list.(n.n)
  in

  let move_related n = node_moves n <> [] in

  let make_worklist () =
    List.iter
      (fun cnode ->
        if degree.(cnode.n) >= k then move_node cnode.setnode spill_worklist
        else if move_related cnode then move_node cnode.setnode freeze_worklist
        else move_node cnode.setnode simplify_worklist)
      (dll2list initial)
  in

  let adjacent cnode =
    Cset.elements
      (Cset.diff adj_list.(cnode.n)
         (Cset.union
            (Cset.of_list (dll2list select_stack))
            (Cset.of_list (dll2list coalesced_nodes))))
  in

  let enable_moves cnodes =
    List.iter
      (fun cnode ->
        List.iter
          (fun m ->
            if m.setnode.dll == active_moves then
              move_node m.setnode worklist_moves)
          (node_moves cnode))
      cnodes
  in

  let decrement_degree cnode =
    let d = degree.(cnode.n) in
    degree.(cnode.n) <- degree.(cnode.n) - 1;
    if d = k then (
      enable_moves (cnode :: adjacent cnode);
      if move_related cnode then move_node cnode.setnode freeze_worklist
      else move_node cnode.setnode simplify_worklist)
  in

  let simplify () =
    List.iter
      (fun (cnode : cnode) ->
        move_node cnode.setnode select_stack;
        List.iter decrement_degree (adjacent cnode))
      (dll2list simplify_worklist)
  in

  let rec get_alias (cnode : cnode) =
    if cnode.setnode.dll == coalesced_nodes then
      get_alias (Option.get alias.(cnode.n))
    else cnode
  in

  let ok t u =
    degree.(t.n) < k || t.setnode.dll == precolored || adj_set.(t.n).(u.n)
  in

  let conservative nodes =
    let d =
      Cset.fold (fun n d -> if degree.(n.n) >= k then d + 1 else d) nodes 0
    in
    d < k
  in

  let add_worklist (u : cnode) =
    if u.setnode.dll != precolored && (not (move_related u)) && degree.(u.n) < k
    then move_node u.setnode simplify_worklist
  in

  let combine (u : cnode) (v : cnode) =
    move_node v.setnode coalesced_nodes;
    alias.(v.n) <- Some u;
    move_list.(u.n) <- move_list.(u.n) @ move_list.(v.n);
    enable_moves [ v ];
    List.iter
      (fun t ->
        add_edge t u;
        decrement_degree t)
      (adjacent v);
    if degree.(u.n) >= k && u.setnode.dll == freeze_worklist then
      move_node u.setnode spill_worklist
  in

  let coalesce () =
    let (({ src; dst } as m) :: _) = dll2list worklist_moves in
    let x = get_alias src in
    let y = get_alias dst in
    let u, v = if y.setnode.dll == precolored then (y, x) else (x, y) in
    if u == v then (
      move_node m.setnode coalesced_moves;
      add_worklist u)
    else if v.setnode.dll == precolored || adj_set.(u.n).(v.n) then (
      move_node m.setnode constrained_moves;
      add_worklist u;
      add_worklist v)
    else if
      u.setnode.dll == precolored
      && List.for_all (fun t -> ok t u) (adjacent v)
      || u.setnode.dll != precolored
         && conservative
              (Cset.union
                 (Cset.of_list (adjacent u))
                 (Cset.of_list (adjacent v)))
    then (
      move_node m.setnode coalesced_moves;
      combine u v;
      add_worklist u)
    else move_node m.setnode active_moves
  in

  let freeze_moves u =
    List.iter
      (fun move ->
        let v =
          if get_alias u == get_alias move.dst then get_alias move.src
          else get_alias move.dst
        in
        move_node move.setnode frozen_moves;
        if (not (move_related v)) && degree.(v.n) < k then
          move_node v.setnode simplify_worklist)
      (node_moves u)
  in

  let freeze () =
    let (u :: _) = dll2list freeze_worklist in
    move_node u.setnode simplify_worklist;
    freeze_moves u
  in

  let select_spill () =
    let cost cnode = spill_cost cnode.inode in
    let (m :: _) =
      List.sort (fun u v -> cost u - cost v) (dll2list spill_worklist)
    in
    move_node m.setnode simplify_worklist;
    freeze_moves m
  in

  let assign_colors () =
    List.iter
      (fun cnode ->
        let ok_colors' =
          Cset.fold
            (fun adj ok_colors ->
              let adj' = get_alias adj in
              if
                adj'.setnode.dll == colored_nodes
                || adj'.setnode.dll == precolored
              then Rset.remove (Option.get color.(adj'.n)) ok_colors
              else ok_colors)
            adj_list.(cnode.n) (Rset.of_list registers)
        in
        if Rset.is_empty ok_colors' then move_node cnode.setnode spilled_nodes
        else (
          move_node cnode.setnode colored_nodes;
          color.(cnode.n) <- Some (Rset.choose ok_colors')))
      (List.rev (dll2list select_stack));
    List.iter
      (fun cnode ->
        let alias = get_alias cnode in
        color.(cnode.n) <- color.(alias.n))
      (dll2list coalesced_nodes)
  in

  build ();
  make_worklist ();

  while
    dll2list simplify_worklist <> []
    || dll2list worklist_moves <> []
    || dll2list freeze_worklist <> []
    || dll2list spill_worklist <> []
  do
    if dll2list simplify_worklist <> [] then simplify ()
    else if dll2list worklist_moves <> [] then coalesce ()
    else if dll2list freeze_worklist <> [] then freeze ()
    else if dll2list spill_worklist <> [] then select_spill ()
  done;

  assign_colors ();

  let allocation' =
    List.fold_left
      (fun alloc cnode ->
        let color = color.(cnode.n) in
        match color with
        | Some c ->
            let t = gtemp cnode.inode in
            Temp.Table.add t c alloc
        | None -> allocation)
      allocation
      (dll2list colored_nodes @ dll2list coalesced_nodes)
  in
  let spilled =
    List.map (fun cnode -> gtemp cnode.inode) (dll2list spilled_nodes)
  in
  (allocation', spilled)
